package org.darchitect.nn.som.ui;

import java.awt.BorderLayout;
import java.awt.Dimension;
import java.awt.FlowLayout;
import java.awt.Menu;
import java.awt.MenuBar;
import java.awt.MenuItem;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.MouseEvent;
import java.awt.event.MouseListener;
import java.awt.geom.Point2D;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Scanner;
import java.util.Stack;
import java.util.StringTokenizer;

import javax.swing.Box;
import javax.swing.JFrame;
import javax.swing.JLabel;
import javax.swing.JOptionPane;
import javax.swing.JPanel;

import org.darchitect.ui.UIUtil;
import org.darchitect.nn.som.SOM;
import org.darchitect.nn.som.SOMController;

public class SOMMainWindow extends JFrame {
	
	private static final long serialVersionUID = 1L;
	private SOMColorGrid visualPanel;
	private JLabel statusLbl = new JLabel(" ");
	private CreateNewSomDialog createSomDialog;
	private QuerySOMDialog queryDialog;
	
	private enum Task { Creating, Training, Querying, Validating; }
	private Stack<Task> tasks = new Stack<Task>();
	
	private SOM som;
	
	public SOMMainWindow() {
		super("Self Organized Maps");
		
		setUpUI();
		
		Dimension dim = new Dimension(600, 500);
		setSize(dim);
		setDefaultCloseOperation(EXIT_ON_CLOSE);
		UIUtil.centerWindow(this);
		UIUtil.setNativeLookAndFeel();
		setVisible(true);
	}
	
	private MouseListener mouseListener = new MouseListener() {
		public void mouseClicked(MouseEvent e) {
			Point2D clickedPoint = e.getPoint();
			visualPanel.selectNodeAt(clickedPoint);
			Point2D p = visualPanel.getNodeCoordinates(clickedPoint);
			double[] w = som.getWeights()[(int)p.getX()][(int)p.getY()];
			DecimalFormat df = new DecimalFormat("0.000");
			String str = "[";
			for (int i = 0; i < w.length; i++) {
				str += df.format(w[i]);
				if (i != w.length - 1)
					str += ", ";
			}
			str += "]";
			setStatus(str);
		}
		public void mouseEntered(MouseEvent e) {
		}
		public void mouseExited(MouseEvent e) {
		}
		public void mousePressed(MouseEvent e) {
		}
		public void mouseReleased(MouseEvent e) {
		}
	};
	
	private BorderLayout bLayout = new BorderLayout();
	
	private void showColorGrid(double w[][][]) {
		SOMColorGrid panel = new SOMColorGrid(w);
		getContentPane().removeAll();
		getContentPane().setLayout(bLayout);
		visualPanel = panel;
		visualPanel.addMouseListener(mouseListener);
		Box vBox = Box.createVerticalBox();
		Box hBox = Box.createHorizontalBox();
		vBox.add(visualPanel);
		JPanel statusPnl = new JPanel();
		statusPnl.setLayout(new BorderLayout());
		statusPnl.add(statusLbl);
		hBox.add(statusPnl);
		vBox.add(hBox);
		getContentPane().add(vBox);
		pack();
		UIUtil.centerWindow(this);
	}
	
	private void setStatus(String msg) {
		statusLbl.setText(msg);
	}
	
	private void openWeights() {
		final File f = UIUtil.chooseFile(UIUtil.OPEN);
		if (f != null) {
			
			UIUtil.ExecutableAndMessageShowable exec = new UIUtil.ExecutableAndMessageShowable() {
				public void execute() {
					SOMController somController = new SOMController();
					som = somController.loadSOMFromWeightFile(f.getAbsolutePath());
					showColorGrid(som.getWeights());
				}
			};
			UIUtil.executeAndShowMessage(this, "We are loading your SOM's weights. One moment please.", 
					"Loading...", exec);
			
		}
	}
	
	private void saveWeights() {
		final File f = UIUtil.chooseFile(UIUtil.SAVE);
		if (f != null) {
			
			UIUtil.ExecutableAndMessageShowable exec = new UIUtil.ExecutableAndMessageShowable() {
				public void execute() {
					SOMController somController = new SOMController();
					somController.saveSOMWeightsInFile(som, f.getAbsolutePath());
				}
			};
			UIUtil.executeAndShowMessage(this, "We are saving your SOM's weights. One moment please.", 
					"Saving...", exec);
			
		}
	}
	
	private void createNewSOMPropertiesPanel() {
		final SOMController somController = new SOMController();
		
		ActionListener createNewSomListener = new ActionListener() {
			public void actionPerformed(ActionEvent e) {
				int layerWidth = createSomDialog.getKohonenLayerWidth();
				int vectorSize = createSomDialog.getInputVectorSize();
				double d0tFactor = createSomDialog.getD0tFactor();
				double a0t = createSomDialog.getA0t();
				int iter = createSomDialog.getIterations();
				createNewSom(layerWidth, vectorSize, d0tFactor, a0t, iter);
				createSomDialog.dispose();
				
				if (! tasks.empty() && tasks.peek() == Task.Training) {
					tasks.pop();
					try {
						trainSOMPanel();
					} catch (Exception e1) {
						JOptionPane.showMessageDialog(SOMMainWindow.this, e1.getMessage(), "Error", JOptionPane.ERROR_MESSAGE);
					}
				}
			}
		};
		
		createSomDialog = new CreateNewSomDialog(somController.getDefaultKohonenLayerWidth(), 
				somController.getDefaultInputVectorSize(), somController.getDefaultD0tFactor(), 
				somController.getDefaultA0t(), somController.getDefaultIterations(), createNewSomListener);
		createSomDialog.setVisible(true);
		
	}
	
	private void querySOMPanel() {
		if (som == null || ! som.hasBeenTrained()) {
			int ans = JOptionPane.showConfirmDialog(this, "A SOM must be trained before querying it.\n" +
					"Should we proceed to train it?", "Info", JOptionPane.YES_NO_OPTION, JOptionPane.QUESTION_MESSAGE);
			if (ans == JOptionPane.YES_OPTION)
				try {
					tasks.push(Task.Querying);
					trainSOMPanel();
				} catch (Exception e1) {
					JOptionPane.showMessageDialog(SOMMainWindow.this, e1.getMessage(), "Error", JOptionPane.ERROR_MESSAGE);
				}
			return;
		}
		ActionListener querySOMListener = new ActionListener() {
			public void actionPerformed(ActionEvent e) {
				double[] qVector = queryDialog.getQueryVector();
				if (qVector == null) {
					JOptionPane.showMessageDialog(SOMMainWindow.this, "Some field values are invalid", 
							"Error", JOptionPane.ERROR_MESSAGE);
					return;
				}
				double[] ans = som.query(qVector);
				queryDialog.setQueryAnswer(ans);
			}
		};
		
		queryDialog = new QuerySOMDialog(som.getWeights()[0][0].length, querySOMListener);
		queryDialog.setVisible(true);
	}
	
	private void validateSOMPanel() {
		
	}
	
	private void trainSOMPanel() throws Exception {
		if (som == null) {
			int ans = JOptionPane.showConfirmDialog(this, "A SOM must be created first in order to train it.\n" +
					"Should we proceed to create one?", "Info", JOptionPane.YES_NO_OPTION, JOptionPane.QUESTION_MESSAGE);
			if (ans == JOptionPane.YES_OPTION) {
				tasks.push(Task.Training);
				createNewSOMPropertiesPanel();
			}
			return;
		}
		try {
			File f = UIUtil.chooseFile(UIUtil.OPEN);
			if (f == null)
				throw new Exception("A training data file must be chosen.");
			Scanner scan = new Scanner(f);
			try {
				int vectorLen = som.getWeights()[0][0].length;
				ArrayList<double[]> td = new ArrayList<double[]>();
				double[] arr;
				while (scan.hasNext()) {
					arr = new double[som.getWeights()[0][0].length];
					StringTokenizer tok;
					tok = new StringTokenizer(scan.nextLine());
					if (tok.countTokens() != vectorLen)
						throw new Exception("Invalid SOM training data");
					for (int j = 0; j < vectorLen; j++)
						arr[j] = Double.parseDouble(tok.nextToken());
					td.add(arr);
				}
				final double[][] w = new double[td.size()][td.get(0).length];
				for (int i = 0; i < w.length; i++)
					w[i] = td.get(i);
				
				UIUtil.ExecutableAndMessageShowable exec = new UIUtil.ExecutableAndMessageShowable() {
					public void execute() {
						som.train(w);
					}
				};
				UIUtil.executeAndShowMessage(this, "We are training hard your SOM. One moment please.", 
						"Training...", exec);
				
				showColorGrid(som.getWeights());
				
				if (! tasks.empty()) {
					Task t = tasks.pop();
					if (t == Task.Querying)
						querySOMPanel();
					else if (t == Task.Validating)
						validateSOMPanel();
				}
			} catch (IOException e) {
				throw new Exception("Error reading training data file");
			}
		} catch (FileNotFoundException e) {
			throw new Exception("File not found?");
		}
	}
	
	private void createNewSom(int kohonenLayerWidth, int inputVectorSize,
			double d0tFactor, double a0t, int iterations) {
		SOMController somController = new SOMController();
		showColorGrid(
				(som = 
					somController.createSOMInstance(kohonenLayerWidth, inputVectorSize, d0tFactor, a0t, iterations)
					).getWeights());
	}

	private void setUpUI() {
		MenuBar bar = new MenuBar();
		Menu fileMenu = new Menu("File");
		MenuItem open = new MenuItem("Load SOM weights...");
		MenuItem save = new MenuItem("Save SOM weights...");
		MenuItem exit = new MenuItem("Exit");
		Menu somMenu = new Menu("SOM");
		MenuItem createSom = new MenuItem("Create new SOM");
		MenuItem trainSom = new MenuItem("Train");
		MenuItem querySom = new MenuItem("Query");
		MenuItem validateSom = new MenuItem("Validate");
		
		open.addActionListener(menuListener);
		save.addActionListener(menuListener);
		exit.addActionListener(menuListener);
		createSom.addActionListener(menuListener);
		trainSom.addActionListener(menuListener);
		querySom.addActionListener(menuListener);
		validateSom.addActionListener(menuListener);
		
		fileMenu.add(open);
		fileMenu.add(save);
		fileMenu.addSeparator();
		fileMenu.add(exit);
		somMenu.add(createSom);
		somMenu.add(trainSom);
		somMenu.add(querySom);
		somMenu.add(validateSom);
		bar.add(fileMenu);
		bar.add(somMenu);
		setMenuBar(bar);
		
		JLabel noSomLbl = new JLabel("Create a new SOM or load its weights");
		setLayout(new FlowLayout());
		getContentPane().add(noSomLbl);
	}
	
	private ActionListener menuListener = new ActionListener() {
		public void actionPerformed(ActionEvent e) {
			MenuItem item = (MenuItem)e.getSource();
			if (item.getLabel().equals("Exit"))
				SOMMainWindow.this.dispose();
			else if (item.getLabel().equals("Load SOM weights..."))
				openWeights();
			else if (item.getLabel().equals("Save SOM weights..."))
				saveWeights();
			else if (item.getLabel().equals("Create new SOM"))
				createNewSOMPropertiesPanel();
			else if (item.getLabel().equals("Train"))
				try {
					trainSOMPanel();
				} catch (Exception e1) {
					JOptionPane.showMessageDialog(SOMMainWindow.this, e1.getMessage(), "Error", JOptionPane.ERROR_MESSAGE);
				}
			else if (item.getLabel().equals("Query"))
				querySOMPanel();
			else if (item.getLabel().equals("Validate"))
				validateSOMPanel();
		}
	};

}
